import torch
import torch.nn as nn
import torch.nn.functional as F


class DepthwiseConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, depth_multiplier=1, padding=0, bias=False):
        super(DepthwiseConv2d, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels * depth_multiplier, kernel_size=kernel_size,
                                   padding=padding, groups=in_channels, bias=bias)
        self.pointwise = nn.Conv2d(in_channels * depth_multiplier, out_channels, kernel_size=1, bias=bias)

    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x


class EEGNetPytorch(nn.Module):
    def __init__(self, nb_classes=2, Chans=64, Samples=128, dropoutRate=0.5, kernLength=64, F1=8, D=2, F2=16,
                 device=None):
        super(EEGNetPytorch, self).__init__()
        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.firstConv = nn.Sequential(
            nn.Conv2d(1, F1, (1, kernLength), padding=(0, kernLength // 2), bias=False),
            nn.BatchNorm2d(F1),
        )
        self.depthwiseConv = DepthwiseConv2d(F1, F1 * D, (Chans, 1), depth_multiplier=D, bias=False)
        self.depthwiseBN = nn.Sequential(
            nn.BatchNorm2d(F1 * D),
            nn.ELU(),
            nn.AvgPool2d((1, 4)),
            nn.Dropout(dropoutRate)
        )
        self.separableConv = DepthwiseConv2d(F1 * D, F2, (1, 16), padding=(0, 8), bias=False)
        self.separableBN = nn.Sequential(
            nn.BatchNorm2d(F2),
            nn.ELU(),
            nn.AvgPool2d((1, 8)),
            nn.Dropout(dropoutRate)
        )
        self.classifier = None  # 不在初始化时创建分类器
        self.nb_classes = nb_classes

    def forward(self, x, return_features=False):
        x = self.firstConv(x)
        x = self.depthwiseConv(x)
        x = self.depthwiseBN(x)
        x = self.separableConv(x)
        x = self.separableBN(x)
        x = x.view(x.size(0), -1)  # Flatten the output
        if self.classifier is None:
            # 动态创建分类器
            num_features = x.shape[1]
            self.classifier = nn.Sequential(
                nn.Linear(num_features, self.nb_classes),
                nn.Softmax(dim=1)
            ).to(self.device)

        x = self.classifier(x)
        if return_features:
            return x  # 返回特征供进一步分析
        return x
